# genetic_algorithm.py（融合CPO的核心实现）
import random
import numpy as np
from typing import List
from vrp_model import VRPTWInstance
from itertools import chain


class EnhancedGeneticAlgorithm:
    def __init__(self, instance: VRPTWInstance,
                 pop_size=100,
                 min_pop_size=50,
                 elite_rate=0.1,
                 crossover_rate=0.8,
                 mutation_rate=0.1,
                 max_generations=200,
                 T=10):
        self.instance = instance
        self.pop_size = pop_size
        self.min_pop_size = min_pop_size
        self.elite_rate = elite_rate
        self.crossover_rate = crossover_rate
        self.mutation_rate = mutation_rate
        self.max_generations = max_generations
        self.T = T
        self.current_pop_size = pop_size
        self.best_solution = None

    def initialize_population(self) -> List[List[int]]:
        """初始化种群"""
        customers = [c.id for c in self.instance.customers if c.id != 0]
        return [random.sample(customers, len(customers)) for _ in range(self.pop_size)]

    def cyclic_population_adjustment(self, gen):
        """CPO动态种群调整"""
        cycle_factor = 1 - (gen % self.T) / self.max_generations
        new_size = int(self.min_pop_size + (self.pop_size - self.min_pop_size) * cycle_factor)
        self.current_pop_size = max(new_size, self.min_pop_size)

    def calculate_fitness(self, chromo):
        """计算适应度（带缓存）"""
        key = tuple(chromo)
        if key in self.instance._fitness_cache:
            return self.instance._fitness_cache[key]

        total_cost = 0
        vehicle_used = 0
        current_route = [0]  # 仓库为0
        current_load = 0
        current_volume = 0
        current_time = 0
        total_distance = 0

        for customer_id in chromo:
            customer = self.instance.customers[customer_id]

            # 检查是否需要新车辆
            if (current_load + customer.demand > self.instance.vehicle.capacity or
                    current_volume + customer.volume > self.instance.vehicle.volume):
                # 计算返回仓库的路径
                return_distance = self.instance.distance_matrix[current_route[-1]][0]
                total_distance += return_distance
                total_cost += return_distance * self.instance.vehicle.distance_cost

                # 车辆使用成本
                total_cost += self.instance.vehicle.fixed_cost
                vehicle_used += 1

                # 重置状态
                current_route = [0]
                current_load = 0
                current_volume = 0
                current_time = 0

            # 计算到当前客户的行驶数据
            from_node = current_route[-1]
            distance = self.instance.distance_matrix[from_node][customer_id]
            travel_time = self.instance.time_matrix[from_node][customer_id]

            # 到达时间和时间窗处理
            arrival_time = current_time + travel_time
            wait_time = max(customer.early - arrival_time, 0)
            time_violation = max(arrival_time - customer.late, 0)

            # 成本累计
            total_distance += distance
            total_cost += distance * self.instance.vehicle.distance_cost  # 距离成本
            total_cost += (travel_time + wait_time) * self.instance.vehicle.time_cost  # 时间成本
            total_cost += time_violation * self.instance.vehicle.time_penalty  # 时间窗违反惩罚

            # 更新状态
            current_route.append(customer_id)
            current_load += customer.demand
            current_volume += customer.volume
            current_time = arrival_time + wait_time + 60  # 服务时间60分钟

            # 最后返回仓库
        return_distance = self.instance.distance_matrix[current_route[-1]][0]
        total_cost += return_distance * self.instance.vehicle.distance_cost
        total_cost += self.instance.vehicle.fixed_cost  # 最后一辆车费用
        vehicle_used += 1

        # 超载惩罚（每次换车时检查）
        overload_penalty = 0
        if current_load > self.instance.vehicle.capacity:
            overload_penalty += self.instance.vehicle.overload_penalty
        if current_volume > self.instance.vehicle.volume:
            overload_penalty += self.instance.vehicle.overload_penalty

        total_cost += overload_penalty

        self.instance._fitness_cache[key] = total_cost
        return total_cost

    def cpo_selection(self, population, fitness):
        """混合选择机制"""
        elite_size = int(self.elite_rate * self.current_pop_size)
        elite_indices = np.argsort(fitness)[:elite_size]
        elites = [population[i] for i in elite_indices]

        # 轮盘赌选择
        inv_fitness = [1 / (f + 1e-6) for f in fitness]
        total = sum(inv_fitness)
        probs = [f / total for f in inv_fitness]

        selected = []
        while len(selected) < self.current_pop_size - elite_size:
            idx = random.choices(range(len(population)), weights=probs)[0]
            parent = population[idx]
            if self.defense_strategy_enhancement(parent, population):
                selected.append(parent)

        return elites + selected

    def defense_strategy_enhancement(self, chromo, population):
        """CPO防御策略增强"""
        current_fitness = self.calculate_fitness(chromo)
        median_fitness = np.median([self.calculate_fitness(c) for c in population])

        if current_fitness > median_fitness:
            return self.exploration_phase(chromo)
        else:
            return self.exploitation_phase(chromo)

    def exploration_phase(self, chromo):
        """探索策略"""
        new_chromo = chromo.copy()
        # 大范围变异
        if len(new_chromo) >= 2:
            i, j = random.sample(range(len(new_chromo)), 2)
            new_chromo[i], new_chromo[j] = new_chromo[j], new_chromo[i]
        # 插入变异
        if random.random() < 0.3:
            gene = new_chromo.pop(random.randint(0, len(new_chromo) - 1))
            new_chromo.insert(random.randint(0, len(new_chromo)), gene)

        if self.calculate_fitness(new_chromo) < self.calculate_fitness(chromo):
            chromo[:] = new_chromo
            return True
        return False

    def exploitation_phase(self, chromo):
        """开发策略"""
        best_chromo = chromo.copy()
        best_fitness = self.calculate_fitness(chromo)
        improved = False

        for _ in range(3):
            new_chromo = chromo.copy()
            pos = random.randint(0, len(new_chromo) - 2)
            new_chromo[pos], new_chromo[pos + 1] = new_chromo[pos + 1], new_chromo[pos]
            new_fitness = self.calculate_fitness(new_chromo)
            if new_fitness < best_fitness:
                best_chromo = new_chromo
                best_fitness = new_fitness
                improved = True

        if improved:
            chromo[:] = best_chromo
        return improved

    def pmx_crossover(self, parent1, parent2):
        """PMX交叉"""
        size = len(parent1)
        p1, p2 = sorted(random.sample(range(size), 2))

        child1 = parent1.copy()
        child2 = parent2.copy()
        mapping1, mapping2 = {}, {}

        for i in range(p1, p2 + 1):
            mapping1[parent2[i]] = parent1[i]
            mapping2[parent1[i]] = parent2[i]
            child1[i], child2[i] = parent2[i], parent1[i]

        for i in chain(range(0, p1), range(p2 + 1, size)):
            while child1[i] in mapping1:
                child1[i] = mapping1[child1[i]]
            while child2[i] in mapping2:
                child2[i] = mapping2[child2[i]]
        print("子代1：{}".format(child1))
        print("子代2：{}".format(child2))
        return child1, child2

    def adaptive_mutation(self, chromo):
        """自适应变异"""
        diversity = len(set(chromo)) / len(chromo)
        rate = self.mutation_rate * (3 if diversity < 0.1 else 2 if diversity < 0.3 else 1)
        return self.mutate(chromo, rate)

    def mutate(self, chromo, rate):
        """基础变异---交换变异"""
        if random.random() < rate:
            i, j = random.sample(range(len(chromo)), 2)
            chromo[i], chromo[j] = chromo[j], chromo[i]
        return chromo

    def solve(self):
        """主求解循环"""
        population = self.initialize_population()
        best_solution = None
        best_fitness = float('inf')
        history = []

        for gen in range(self.max_generations):
            self.cyclic_population_adjustment(gen)
            population = population[:self.current_pop_size]

            fitness = [self.calculate_fitness(c) for c in population]
            selected = self.cpo_selection(population, fitness)

            new_population = []
            while len(new_population) < self.current_pop_size:
                parents = random.sample(selected, 2)
                if random.random() < self.crossover_rate:
                    child1, child2 = self.pmx_crossover(parents[0], parents[1])
                else:
                    child1, child2 = parents[0].copy(), parents[1].copy()

                new_population.extend([
                    self.adaptive_mutation(child1),
                    self.adaptive_mutation(child2)
                ])

            population = new_population[:self.current_pop_size]
            current_best = min(population, key=self.calculate_fitness)
            current_fitness = self.calculate_fitness(current_best)

            if current_fitness < best_fitness:
                best_solution = current_best
                best_fitness = current_fitness

            history.append(best_fitness)
            print(f"Gen {gen + 1}: Cost={best_fitness:.2f} Pop={self.current_pop_size}")
            self.instance.clear_cache()  # 每代清除缓存

        return best_solution, history